package cn.seecoder.ai.algorithm.ml.classification

import cn.seecoder.ai.algorithm.ml.BaseClassification
import cn.seecoder.ai.enums.TrainParamEnum
import cn.seecoder.ai.model.bo.TrainParamsBO
import cn.seecoder.ai.utils.TrainParamReader.readParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.LogisticRegression

import java.util

class LogisticRegressionClassification extends BaseClassification{

  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {

    val lr = new LogisticRegression()
      .setFeaturesCol("indexedFeatures")
      .setLabelCol("label")
      .setMaxIter(trainParams.maxIter)
      .setRegParam(trainParams.regParam)
      .setElasticNetParam(trainParams.elasticNetParam)

    lr
  }


}
